import sys
import face_recognition
import os
import multiprocessing
from functools import partial
import json
from typing import List, Dict, Any
import argparse

class FileInfo:
    """存储文件信息的类"""
    def __init__(self, file_name: str, file_path: str):
        self.file_name = file_name
        self.file_path = file_path
        self.people_names: List[str] = []

    def to_dict(self) -> Dict[str, Any]:
        """将FileInfo对象转换为字典"""
        return {
            "fileName": self.file_name,
            "filePath": self.file_path,
            "peopleNames": self.people_names if self.people_names else ["未知"]
        }

def process_image(image_path: str, known_face_encodings: List[Any], known_face_names: List[str], tolerance: float) -> FileInfo:
    """处理单个图像文件，识别人脸并返回FileInfo对象"""
    file_info = FileInfo(
        file_name=os.path.basename(image_path),
        file_path=image_path
    )

    # 加载图像并识别人脸
    image = face_recognition.load_image_file(image_path)
    face_locations = face_recognition.face_locations(image)
    face_encodings = face_recognition.face_encodings(image, face_locations)
    
    if face_encodings:
        for face_encoding in face_encodings:
            # 比较人脸编码
            matches = face_recognition.compare_faces(known_face_encodings, face_encoding, tolerance=tolerance)
            face_distances = face_recognition.face_distance(known_face_encodings, face_encoding)
            
            best_match_index = face_distances.argmin()
            if matches[best_match_index]:
                name = known_face_names[best_match_index]
                if name not in file_info.people_names:
                    file_info.people_names.append(name)
    
    # 如果没有检测到人脸或没有匹配的人脸，people_names 将保持为空列表

    return file_info

def get_image_files(path: str) -> List[str]:
    """递归获取目录中的所有图像文件"""
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif'}
    image_files = []
    if os.path.isfile(path):
        if os.path.splitext(path)[1].lower() in image_extensions:
            image_files.append(path)
    elif os.path.isdir(path):
        for root, _, files in os.walk(path):
            for file in files:
                if os.path.splitext(file)[1].lower() in image_extensions:
                    image_files.append(os.path.join(root, file))
    return image_files

def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description='人脸识别脚本')
    parser.add_argument('--cpus', type=int, default=1, help='使用的CPU核心数')
    parser.add_argument('--tolerance', type=float, default=0.6, help='人脸识别容差')
    parser.add_argument('--show-distance', type=str, choices=['true', 'false'], default='false', help='显示距离（当前版本未使用）')
    parser.add_argument('known_faces_dir', type=str, help='已知人脸目录')
    parser.add_argument('unknown_faces_input', type=str, help='待识别的图像文件或目录')

    args = parser.parse_args()

    # 设置CPU核心数
    cpus = args.cpus if args.cpus > 0 else multiprocessing.cpu_count()
    tolerance = args.tolerance
    known_faces_dir = os.path.expanduser(args.known_faces_dir)
    unknown_faces_input = os.path.expanduser(args.unknown_faces_input)
    
    # 加载已知人脸
    known_face_encodings = []
    known_face_names = []
    for image_name in os.listdir(known_faces_dir):
        image_path = os.path.join(known_faces_dir, image_name)
        image = face_recognition.load_image_file(image_path)
        face_encodings = face_recognition.face_encodings(image)
        if face_encodings:
            known_face_encodings.append(face_encodings[0])
            known_face_names.append(os.path.splitext(image_name)[0])
    
    # 获取待识别图像的路径
    unknown_image_paths = get_image_files(unknown_faces_input)
    
    if not unknown_image_paths:
        print(f"错误：在 '{unknown_faces_input}' 中没有找到有效的图像文件")
        sys.exit(1)
    
    # 使用多进程处理图像
    with multiprocessing.Pool(processes=cpus) as pool:
        process_func = partial(process_image, 
                               known_face_encodings=known_face_encodings, 
                               known_face_names=known_face_names, 
                               tolerance=tolerance)
        results = list(pool.imap(process_func, unknown_image_paths))
    
    # 将结果转换为字典列表
    json_results = [file_info.to_dict() for file_info in results]
    
    # 将JSON结果输出到控制台
    print(json.dumps(json_results, ensure_ascii=False, indent=4))

if __name__ == "__main__":
    main()
